//! ## DBT Routing for Both Clients and Server in CS2 System
//!
//! The routing are used to route:
//! 1. the cache attachment & lookup operations of clients to the corresponding attached
//!    cache tables.
//! 2. the cache creating & synchronization operations of servers to the corresponding
//!    created cache tables.

use std::mem;

use rustc_hash::FxHashMap;

use crate::{
    cache::{CacheTableClientHandle, CacheTableKeyRef, CacheTableServerHandle, CacheTableVal},
    error::{Error, Result},
    utils::mangle::PathStrExt,
};

/// The router for cache attachment & lookup operations of clients.
pub struct LookupRouter {
    prefix: String,
    tables: FxHashMap<String, CacheTableClientHandle>,
}

impl LookupRouter {
    /// Create a new router with the given prefix and permissive flag.
    pub fn new(prefix: &str) -> Self {
        Self {
            prefix: prefix.to_string(),
            tables: FxHashMap::default(),
        }
    }

    /// Attach the given paths to the router.
    pub fn attach<L, S>(&mut self, paths: L) -> Result<()>
    where
        L: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        for path in paths {
            let path = path.as_ref();

            if self.tables.contains_key(path) {
                continue;
            }

            self.tables.insert(
                path.to_string(),
                CacheTableClientHandle::new(&format!("{}-{}", self.prefix, path.mangle()))?,
            );
        }

        Ok(())
    }

    /// Lookup the given guest code in the cache tables according to the given path.
    pub fn lookup<'a>(
        &'a mut self,
        path: impl AsRef<str>,
        cache_key: &CacheTableKeyRef<'a>,
    ) -> Result<Option<CacheTableVal<'a>>> {
        let path = path.as_ref();

        let table = self
            .tables
            .get_mut(path)
            .ok_or_else(|| Error::TableNotFound {
                path: path.to_string(),
            })?;

        loop {
            match table.lookup(cache_key) {
                Err(Error::EnlargementNeeded) => table.enlarge()?,
                // Safety: make the stupid borrow checker happy.
                //
                // Workaround for https://github.com/rust-lang/rust/issues/54663
                result => unsafe {
                    return mem::transmute::<
                        Result<Option<CacheTableVal<'_>>>,
                        Result<Option<CacheTableVal<'_>>>,
                    >(result);
                },
            }
        }
    }

    /// Iterate the cache data in the cache tables according to the given path. The
    /// function enlarges the cache table if needed.
    pub fn iter_data(
        &mut self,
        path: impl AsRef<str>,
    ) -> Result<impl Iterator<Item = (CacheTableKeyRef, CacheTableVal)>> {
        let path = path.as_ref();

        let table = self
            .tables
            .get_mut(path)
            .ok_or_else(|| Error::TableNotFound {
                path: path.to_string(),
            })?;

        table.enlarge()?;

        Ok(table.iter_data())
    }
}

/// The router for cache creating & synchronization operations of servers.
pub struct HandleRouter {
    prefix: String,
    tables: FxHashMap<String, CacheTableServerHandle>,
}

impl HandleRouter {
    /// The default granularity for cache table enlargement.
    const GRANULARITY: usize = 4;

    /// The default limit for cache table enlargement.
    const ENLARGEMENT_LIMIT: usize = 1024;

    /// Create a new router with the given prefix and permissive flag.
    pub fn new(prefix: &str) -> Self {
        Self {
            prefix: prefix.to_string(),
            tables: FxHashMap::default(),
        }
    }

    /// Create cache tables for the given paths.
    pub fn create<L, S>(&mut self, paths: L) -> Result<()>
    where
        L: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        for path in paths {
            let path = path.as_ref();

            if self.tables.contains_key(path) {
                continue;
            }

            self.tables.insert(
                path.to_string(),
                CacheTableServerHandle::new(
                    &format!("{}-{}", self.prefix, path.mangle()),
                    1024,              // TODO: make it configurable
                    Self::GRANULARITY, // TODO: make it configurable
                )?,
            );
        }

        Ok(())
    }

    /// Synchronize the given host meta/code to the cache tables according to the given
    /// path.
    pub fn synchronize(
        &mut self,
        path: impl AsRef<str>,
        cache_key: &CacheTableKeyRef,
        cache_val: &CacheTableVal,
    ) -> Result<()> {
        let path = path.as_ref();

        let table = self
            .tables
            .get_mut(path)
            .ok_or_else(|| Error::TableNotFound {
                path: path.to_string(),
            })?;

        while let Err(Error::MemNotEnough { required, free }) = table.insert(cache_key, cache_val) {
            let enlargement = (table.capacity() * 3 / 2)
                .div_ceil(Self::GRANULARITY * 4096)
                .min(Self::ENLARGEMENT_LIMIT);

            table.enlarge(
                (required - free)
                    .next_multiple_of(Self::GRANULARITY * 4096)
                    .div_ceil(Self::GRANULARITY * 4096)
                    .max(enlargement),
            )?;
        }

        Ok(())
    }
}

#[cfg(test)]
mod test {
    use coverage_helper::test;
    use cs2s_testdata::*;

    use crate::cache::CacheTableKeyRef;

    use super::*;

    #[test]
    fn basic_interaction() {
        const PREFIX: &str = "test-route-basic-interaction";

        let mut lookup_router = LookupRouter::new(PREFIX);
        let mut handle_router = HandleRouter::new(PREFIX);

        // Create & attach cache tables.
        assert!(handle_router.create(PATHS).is_ok());
        assert!(lookup_router.attach(PATHS).is_ok());

        // Assert that the cache tables are empty.
        for path in PATHS {
            for (i, guest_code) in GUEST_CODES.iter().enumerate() {
                assert!(lookup_router
                    .lookup(
                        path,
                        &CacheTableKeyRef {
                            guest_addr: 0x400000 + i * 0x40,
                            guest_size: guest_code.len(),
                            guest_sign: &guest_code.into(),
                        },
                    )
                    .unwrap()
                    .is_none());
            }
        }

        // Synchronize cache & test lookup.
        for (i, (&path, (&guest_code, (&host_meta, &host_code)))) in PATHS
            .iter()
            .zip(
                GUEST_CODES
                    .iter()
                    .zip(HOST_METAS.iter().zip(HOST_CODES.iter())),
            )
            .enumerate()
        {
            let guest_sign = guest_code.into();

            assert!(handle_router
                .synchronize(
                    path,
                    &CacheTableKeyRef {
                        guest_addr: 0x400000 + i * 0x40,
                        guest_size: guest_code.len(),
                        guest_sign: &guest_sign,
                    },
                    &CacheTableVal {
                        host_meta,
                        host_code,
                    },
                )
                .is_ok());

            let data = lookup_router
                .lookup(
                    path,
                    &CacheTableKeyRef {
                        guest_addr: 0x400000 + i * 0x40,
                        guest_size: guest_code.len(),
                        guest_sign: &guest_sign,
                    },
                )
                .unwrap()
                .unwrap();

            assert_eq!(data.host_meta, host_meta);
            assert_eq!(data.host_code, host_code);

            assert!(lookup_router.iter_data(path).unwrap().any(|(key, val)| {
                key == CacheTableKeyRef {
                    guest_addr: 0x400000 + i * 0x40,
                    guest_size: guest_code.len(),
                    guest_sign: &guest_sign,
                } && val.host_meta == host_meta
                    && val.host_code == host_code
            }));
        }
    }

    #[test]
    fn auto_enlargement() {
        const PREFIX: &str = "test-route-auto-enlargement";
        const LARGE_GUEST_CODE: &[u8] = &[0xCC; 4096 * HandleRouter::GRANULARITY];
        const LARGE_HOST_META: &[u8] = &[0x55; 4096 * HandleRouter::GRANULARITY];
        const LARGE_HOST_CODE: &[u8] = &[0x90; 4096 * HandleRouter::GRANULARITY];

        let mut lookup_router = LookupRouter::new(PREFIX);
        let mut handle_router = HandleRouter::new(PREFIX);

        // Create & attach cache tables.
        assert!(handle_router.create(PATHS).is_ok());
        assert!(lookup_router.attach(PATHS).is_ok());

        handle_router
            .synchronize(
                PATH_0,
                &CacheTableKeyRef {
                    guest_addr: 0x400000,
                    guest_size: LARGE_GUEST_CODE.len(),
                    guest_sign: &LARGE_GUEST_CODE.into(),
                },
                &CacheTableVal {
                    host_meta: LARGE_HOST_META,
                    host_code: LARGE_HOST_CODE,
                },
            )
            .unwrap();

        lookup_router
            .lookup(
                PATH_0,
                &CacheTableKeyRef {
                    guest_addr: 0x400000,
                    guest_size: LARGE_GUEST_CODE.len(),
                    guest_sign: &LARGE_GUEST_CODE.into(),
                },
            )
            .unwrap()
            .unwrap();
    }

    #[test]
    fn path_attach_error() {
        const PREFIX: &str = "test-route-duplicated-path-attach";

        let mut lookup_router = LookupRouter::new(PREFIX);
        let mut handle_router = HandleRouter::new(PREFIX);

        // Should not attach the non-existing paths.
        assert!(lookup_router.attach(PATHS).is_err());

        // Create & attach cache tables.
        assert!(handle_router.create(PATHS).is_ok());
        assert!(lookup_router.attach(PATHS).is_ok());

        // Create the paths again.
        assert!(handle_router.create(PATHS).is_ok());

        // Attach the paths again.
        assert!(lookup_router.attach(PATHS).is_ok());
    }

    #[test]
    fn path_not_found() {
        const PREFIX: &str = "test-route-path-not-found";

        let mut lookup_router = LookupRouter::new(PREFIX);
        let mut handle_router = HandleRouter::new(PREFIX);

        // Create & attach cache tables.
        assert!(handle_router.create(PATHS).is_ok());
        assert!(lookup_router.attach(PATHS).is_ok());

        // Synchronize to the non-existing path.
        assert!(handle_router
            .synchronize(
                "non-existing-path",
                &CacheTableKeyRef {
                    guest_addr: 0x400000,
                    guest_size: 0x40,
                    guest_sign: &GUEST_CODES[0].into(),
                },
                &CacheTableVal {
                    host_meta: HOST_METAS[0],
                    host_code: HOST_CODES[0],
                },
            )
            .is_err());

        // Lookup the non-existing path.
        assert!(lookup_router
            .lookup(
                "non-existing-path",
                &CacheTableKeyRef {
                    guest_addr: 0x400000,
                    guest_size: 0x40,
                    guest_sign: &GUEST_CODES[0].into(),
                },
            )
            .is_err());

        // Iterate the cache blocks of the non-existing path.
        assert!(lookup_router.iter_data("non-existing-path").is_err());
    }
}
